import * as vscode from 'vscode';
import * as path from 'path';
import * as fs from 'fs';

// 用于标记函数选中状态
interface SelectedFunction {
  name: string;
  selected: boolean;
}

// 增强文件信息，包含函数选中状态
interface FileWithFunctions {
  path: string;
  functions: SelectedFunction[];
}

class TestGeneratorProvider implements vscode.TreeDataProvider<vscode.TreeItem> {
  private _onDidChangeTreeData: vscode.EventEmitter<vscode.TreeItem | undefined> = new vscode.EventEmitter<vscode.TreeItem | undefined>();
  readonly onDidChangeTreeData: vscode.Event<vscode.TreeItem | undefined> = this._onDidChangeTreeData.event;

  private selectedFiles: FileWithFunctions[] = []; 
  private showFunctions: boolean = false;
  private currentFileIndex: number = -1;

  getTreeItem(element: vscode.TreeItem): vscode.TreeItem {
    return element;
  }

  getChildren(element?: vscode.TreeItem): Thenable<vscode.TreeItem[]> {
    if (!element) {
      const items: vscode.TreeItem[] = [];

      // 1. 选择文件标题
      const titleItem = new vscode.TreeItem('1. 选择文件', vscode.TreeItemCollapsibleState.None);
      items.push(titleItem);

      // 选择文件按钮
      const selectFileItem = new vscode.TreeItem('Select Files', vscode.TreeItemCollapsibleState.None);
      selectFileItem.command = {
        command: 'vscode-plus.selectFile',
        title: 'Select Files'
      };
      items.push(selectFileItem);

      if (this.selectedFiles.length > 0) {
        this.selectedFiles.forEach((file, index) => {
          const fileItem = new vscode.TreeItem(
            `${path.basename(file.path)}`, 
            vscode.TreeItemCollapsibleState.None
          );
          fileItem.iconPath = new vscode.ThemeIcon('file-code');
          fileItem.tooltip = file.path;
          fileItem.command = {
            command: 'vscode-plus.toggleFileFunctions',
            title: 'Toggle Functions',
            arguments: [index]
          };
          
          
          items.push(fileItem);
        });
      }

      // 2. 选择函数标题
      if (this.showFunctions && this.currentFileIndex >= 0) {
        const currentFile = this.selectedFiles[this.currentFileIndex];
        const funcGroup = new vscode.TreeItem(
          `2. 为 "${path.basename(currentFile.path)}" 选择函数 (可选)`, 
          vscode.TreeItemCollapsibleState.None
        );
        items.push(funcGroup);

        // 生成测试按钮
        const selectedCount = currentFile.functions.filter(f => f.selected).length;
        const generateBtn = new vscode.TreeItem(
          `生成测试 (${selectedCount})`, 
          vscode.TreeItemCollapsibleState.None
        );
        generateBtn.iconPath = new vscode.ThemeIcon('play');
        generateBtn.command = {
          command: 'vscode-plus.generateTestForSelected',
          title: 'Generate Test',
          arguments: [this.currentFileIndex]
        };
        items.push(generateBtn);

        // 全选复选框
        const allSelected = currentFile.functions.every(f => f.selected);
        const checkBoxState = allSelected ? vscode.TreeItemCheckboxState.Checked : vscode.TreeItemCheckboxState.Unchecked;
        const selectAllItem = new vscode.TreeItem('全选', vscode.TreeItemCollapsibleState.None);
        selectAllItem.checkboxState = checkBoxState;
        selectAllItem.command = {
          command: 'vscode-plus.selectAllFunctions',
          title: 'Select All',
          arguments: [this.currentFileIndex]
        };
        items.push(selectAllItem);

        // 函数列表
        currentFile.functions.forEach(func => {
          const funcItem = new vscode.TreeItem(func.name, vscode.TreeItemCollapsibleState.None);
          funcItem.iconPath = new vscode.ThemeIcon('symbol-function');
          funcItem.checkboxState = func.selected ? vscode.TreeItemCheckboxState.Checked : vscode.TreeItemCheckboxState.Unchecked;
          funcItem.command = {
            command: 'vscode-plus.toggleFunction',
            title: 'Toggle Function',
            arguments: [this.currentFileIndex, func.name]
          };
          items.push(funcItem);
        });
      }

      return Promise.resolve(items);
    }

    return Promise.resolve([]);
  }

  public addSelectedFiles(filePaths: string[], functionsList: string[][]) {
    const newFiles = filePaths.map((path, index) => ({
      path,
      functions: functionsList[index].map(name => ({ name, selected: false }))
    }));
    
    // 过滤重复文件
    const existingPaths = this.selectedFiles.map(f => f.path);
    const uniqueNewFiles = newFiles.filter(f => !existingPaths.includes(f.path));
    
    this.selectedFiles = [...this.selectedFiles, ...uniqueNewFiles];
    this._onDidChangeTreeData.fire(undefined);
  }

  // 移除文件
  public removeFile(index: number) {
    this.selectedFiles.splice(index, 1);
    // 如果移除的是当前正在查看的文件，重置状态
    if (index === this.currentFileIndex) {
      this.currentFileIndex = -1;
      this.showFunctions = false;
    } else if (index < this.currentFileIndex) {
      // 如果移除的是当前文件之前的文件，调整索引
      this.currentFileIndex--;
    }
    this._onDidChangeTreeData.fire(undefined);
  }

  // 添加获取已选文件数量的方法
  public getSelectedFilesCount(): number {
    return this.selectedFiles.length;
  }

  public toggleFileFunctions(fileIndex: number) {
    this.currentFileIndex = this.currentFileIndex === fileIndex ? -1 : fileIndex;
    this.showFunctions = this.currentFileIndex !== -1;
    this._onDidChangeTreeData.fire(undefined);
  }

  public selectAllFunctions(fileIndex: number) {
    const file = this.selectedFiles[fileIndex];
    if (file) {
      const allSelected = file.functions.every(f => f.selected);
      file.functions.forEach(f => f.selected = !allSelected);
      this._onDidChangeTreeData.fire(undefined);
    }
  }

  public toggleFunction(fileIndex: number, funcName: string) {
    const file = this.selectedFiles[fileIndex];
    if (file) {
      const func = file.functions.find(f => f.name === funcName);
      if (func) {
        func.selected = !func.selected;
        this._onDidChangeTreeData.fire(undefined);
      }
    }
  }

  public clearSelection() {
    this.selectedFiles = [];
    this.currentFileIndex = -1;
    this.showFunctions = false;
    this._onDidChangeTreeData.fire(undefined);
  }

  public getFileInfo(index: number): FileWithFunctions | undefined {
    if (index >= 0 && index < this.selectedFiles.length) {
      return this.selectedFiles[index];
    }
    return undefined;
  }
}

export function activate(context: vscode.ExtensionContext) {
  const testGeneratorProvider = new TestGeneratorProvider();
  context.subscriptions.push(
    vscode.window.registerTreeDataProvider('testGeneratorView', testGeneratorProvider)
  );

  let panel: vscode.WebviewPanel | undefined = undefined;

  const watcher = vscode.workspace.createFileSystemWatcher('**/*.{cpp,h,hpp}');

  watcher.onDidChange(async uri => {
    const newPanel = getOrCreatePanel(context);
    if (!newPanel) {
      vscode.window.showErrorMessage('无法创建Webview面板');
      return;
    }
    panel = newPanel;
    try {
      panel.reveal();
    } catch (error) {
      vscode.window.showErrorMessage(`显示面板失败: ${error}`);
    }
  });

  // 文件选择命令（改为追加模式）
  context.subscriptions.push(
    vscode.commands.registerCommand('vscode-plus.selectFile', async () => {
      const workspaceFolders = vscode.workspace.workspaceFolders;
      if (!workspaceFolders || workspaceFolders.length === 0) {
        vscode.window.showErrorMessage('请先打开一个工作区');
        return;
      }

      const fileUris = await vscode.window.showOpenDialog({
        canSelectFiles: true,
        canSelectFolders: false,
        canSelectMany: true,
        filters: {
          'C/C++ Files': ['c', 'cpp', 'cc', 'cxx', 'h', 'hpp', 'hh', 'ixx', 'cppm']
        },
        defaultUri: workspaceFolders[0].uri
      });

      if (fileUris && fileUris.length > 0) {
        const invalidFiles = fileUris.filter(uri => 
          !workspaceFolders.some(folder => uri.fsPath.startsWith(folder.uri.fsPath))
        );
        
        if (invalidFiles.length > 0) {
          vscode.window.showErrorMessage('请选择工作区内的文件');
          return;
        }

        try {
          const fileContents = await Promise.all(
            fileUris.map(uri => fs.promises.readFile(uri.fsPath, 'utf-8'))
          );
          const functionsList = fileContents.map(content => extractFunctions(content));
          
          testGeneratorProvider.addSelectedFiles(
            fileUris.map(uri => uri.fsPath),
            functionsList
          );
          
          const newPanel = getOrCreatePanel(context);
          if (newPanel) {
            panel = newPanel;
            panel.reveal();

            // 判断是否是第一次选择文件（之前没有任何文件）
            if (testGeneratorProvider.getSelectedFilesCount() === fileUris.length) {
              testGeneratorProvider.toggleFileFunctions(0);
              const firstFile = testGeneratorProvider.getFileInfo(0);
              if (firstFile) {
                panel.webview.postMessage({
                  command: 'selectFile',
                  filePath: firstFile.path,
                  functions: firstFile.functions.map(f => f.name)
                });
              }
            }
          } else {
            vscode.window.showErrorMessage('创建Webview面板失败');
          }
        } catch (error) {
          vscode.window.showErrorMessage(`读取文件失败: ${error}`);
        }
      }
    })
  );

  // 添加删除文件命令
  context.subscriptions.push(
    vscode.commands.registerCommand('vscode-plus.removeFile', (index: number) => {
      testGeneratorProvider.removeFile(index);
      vscode.window.showInformationMessage('文件已移除');
    })
  );

  // 其他命令保持不变...
  context.subscriptions.push(
    vscode.commands.registerCommand('vscode-plus.toggleFileFunctions', (fileIndex: number) => {
      testGeneratorProvider.toggleFileFunctions(fileIndex);
      
      const fileInfo = testGeneratorProvider.getFileInfo(fileIndex);
      if (fileInfo) {
        const currentPanel = getOrCreatePanel(context);
        if (currentPanel) {
          panel = currentPanel;
          currentPanel.webview.postMessage({
            command: 'selectFile',
            filePath: fileInfo.path,
            functions: fileInfo.functions.map(f => f.name)
          });
        }
      }
    })
  );

  context.subscriptions.push(
    vscode.commands.registerCommand('vscode-plus.selectAllFunctions', (fileIndex: number) => {
      testGeneratorProvider.selectAllFunctions(fileIndex);
    })
  );

  context.subscriptions.push(
    vscode.commands.registerCommand('vscode-plus.toggleFunction', (fileIndex: number, funcName: string) => {
      testGeneratorProvider.toggleFunction(fileIndex, funcName);
    })
  );

  context.subscriptions.push(
    vscode.commands.registerCommand('vscode-plus.generateTestForSelected', async (fileIndex: number) => {
      const fileInfo = testGeneratorProvider.getFileInfo(fileIndex);
      if (!fileInfo) {
        vscode.window.showErrorMessage('未找到文件信息');
        return;
      }

      const selectedFuncs = fileInfo.functions.filter(f => f.selected).map(f => f.name);
      if (selectedFuncs.length === 0) {
        vscode.window.showInformationMessage('请选择至少一个函数');
        return;
      }

      try {
        const sourceCode = await fs.promises.readFile(fileInfo.path, 'utf-8');
        const testFilePath = getTestFilePath(fileInfo.path);
        let testContent = `// 自动生成的单元测试文件 for ${path.basename(fileInfo.path)}\n`;
        testContent += '#include <gtest/gtest.h>\n';
        const includes = extractIncludes(sourceCode);
        testContent += includes.map(inc => `#include "${inc}"`).join('\n');
        testContent += '\n\n';

        const className = extractClassName(sourceCode);
        const classTestPrefix = className ? `class ${className}Test : public ::testing::Test {\nprotected:\n  ${className} obj;\n};\n\n` : '';
        testContent += classTestPrefix;

        selectedFuncs.forEach(funcName => {
          if (className) {
            testContent += `TEST_F(${className}Test, ${funcName}Test) {\n`;
            testContent += `  // 测试 ${className}::${funcName} 方法\n`;
            testContent += `  // TODO: 添加测试逻辑\n`;
            testContent += `  EXPECT_EQ(obj.${funcName}(), 0);\n`;
            testContent += `}\n\n`;
          } else {
            testContent += `TEST(${funcName}Test, BasicTest) {\n`;
            testContent += `  // 测试 ${funcName} 函数\n`;
            testContent += `  // TODO: 添加测试逻辑\n`;
            testContent += `  EXPECT_EQ(${funcName}(), 0);\n`;
            testContent += `}\n\n`;
          }
        });

        await fs.promises.writeFile(testFilePath, testContent, 'utf-8');
        const doc = await vscode.workspace.openTextDocument(testFilePath);
        await vscode.window.showTextDocument(doc);
        vscode.window.showInformationMessage(`成功生成测试文件: ${path.basename(testFilePath)}`);
      } catch (error) {
        vscode.window.showErrorMessage(`生成测试失败: ${error}`);
      }
    })
  );

  context.subscriptions.push(
    vscode.commands.registerCommand('vscode-plus.openTestGenerator', () => {
      const newPanel = getOrCreatePanel(context);
      if (newPanel) {
        panel = newPanel;
        newPanel.reveal();
      } else {
        vscode.window.showErrorMessage('创建Webview面板失败');
      }
    })
  );

  context.subscriptions.push(
    vscode.commands.registerCommand('vscode-plus.clearSelection', () => {
      testGeneratorProvider.clearSelection();
      vscode.window.showInformationMessage('已清除文件选择');
    })
  );
}

// 工具函数保持不变...
function getOrCreatePanel(context: vscode.ExtensionContext): vscode.WebviewPanel | undefined {
  try {
    let panel: vscode.WebviewPanel | undefined;
    if (!panel) {
      panel = vscode.window.createWebviewPanel(
        'testGenerator',
        'C++ Unit Test Generator',
        vscode.ViewColumn.One,
        {
          enableScripts: true,
          localResourceRoots: [
            vscode.Uri.file(path.join(context.extensionPath, 'src/test-generator/build'))
          ]
        }
      );
      const indexPath = vscode.Uri.file(
        path.join(context.extensionPath, 'src/test-generator/build/index.html')
      );
      panel.webview.html = getWebviewContent(indexPath, panel.webview, context.extensionPath);
      panel.onDidDispose(() => {
        panel = undefined;
      });
    }
    return panel;
  } catch (error) {
    vscode.window.showErrorMessage(`创建Webview面板失败: ${error}`);
    return undefined;
  }
}

function getWebviewContent(indexUri: vscode.Uri, webview: vscode.Webview, extensionPath: string) {
  const baseUri = webview.asWebviewUri(indexUri);
  return `<!DOCTYPE html>
    <html lang="en">
      <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>Test Generator</title>
        <link rel="stylesheet" href="${webview.asWebviewUri(vscode.Uri.file(
          path.join(extensionPath, 'src/test-generator/dist/main.css')
        ))}">
      </head>
      <body>
        <div id="root"></div>
        <script src="${webview.asWebviewUri(vscode.Uri.file(
          path.join(extensionPath, 'src/test-generator/dist/main.js')
        ))}"></script>
      </body>
    </html>`;
}

function getTestFilePath(sourcePath: string): string {
  const dirName = path.dirname(sourcePath);
  const baseName = path.basename(sourcePath, path.extname(sourcePath));
  return path.join(dirName, `${baseName}_test.cpp`);
}

function generateFunctionTest(sourceCode: string, sourceFileName: string, functionName: string): string {
  const includes = extractIncludes(sourceCode);
  const className = extractClassName(sourceCode);
  const baseName = path.basename(sourceFileName, path.extname(sourceFileName));
  
  let testContent = `// 自动生成的单元测试文件 for ${sourceFileName}\n`;
  testContent += '#include <gtest/gtest.h>\n';
  testContent += includes.map(inc => `#include "${inc}"`).join('\n');
  testContent += '\n\n';
  
  if (className) {
    testContent += `class ${className}Test : public ::testing::Test {\n`;
    testContent += `protected:\n`;
    testContent += `  ${className} obj;\n`;
    testContent += `};\n\n`;
    
    testContent += `TEST_F(${className}Test, ${functionName}Test) {\n`;
    testContent += `  // 测试 ${className}::${functionName} 方法\n`;
    testContent += `  // TODO: 添加测试逻辑\n`;
    testContent += `  EXPECT_EQ(obj.${functionName}(), 0);\n`;
    testContent += `}\n`;
  } else {
    testContent += `TEST(${functionName}Test, BasicTest) {\n`;
    testContent += `  // 测试 ${functionName} 函数\n`;
    testContent += `  // TODO: 添加测试逻辑\n`;
    testContent += `  EXPECT_EQ(${functionName}(), 0);\n`;
    testContent += `}\n`;
  }

  return testContent;
}

function extractIncludes(sourceCode: string): string[] {
  const includeRegex = /#include\s+["<](.*?)[">]/g;
  const includes: string[] = [];
  let match;
  
  while ((match = includeRegex.exec(sourceCode)) !== null) {
    if (!match[1].includes('gtest') && !match[1].includes('gmock')) {
      includes.push(match[1]);
    }
  }

  return includes;
}

function extractFunctions(sourceCode: string): string[] {
  const functionRegex = /\b([a-zA-Z_][a-zA-Z0-9_]*)\s*\([^)]*\)\s*\{/g;
  const functions: string[] = [];
  
  const excludeList = [
    'if', 'else', 'for', 'while', 'do', 'switch', 'case', 'return',
    'break', 'continue', 'goto', 'sizeof', 'typedef', 'struct', 
    'union', 'enum', 'static', 'extern', 'auto', 'register', 
    'const', 'volatile', 'inline', 'restrict', 'asm', 'typeof'
  ];
  
  let match;
  while ((match = functionRegex.exec(sourceCode)) !== null) {
    const funcName = match[1];
    if (!excludeList.includes(funcName)) {
      functions.push(funcName);
    }
  }

  return functions;
}

function extractClassName(sourceCode: string): string | null {
  const classRegex = /class\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*\{/g;
  let match;
  
  if ((match = classRegex.exec(sourceCode)) !== null) {
    return match[1];
  }
  
  const structRegex = /struct\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*\{/g;
  if ((match = structRegex.exec(sourceCode)) !== null) {
    return match[1];
  }
  
  const enumRegex = /enum\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*\{/g;
  if ((match = enumRegex.exec(sourceCode)) !== null) {
    return match[1];
  }
  
  return null;
}

export function deactivate() { }